/*
 * Copyright (c) 2022. China Mobile (SuZhou) Software Technology Co.,Ltd. All rights reserved.
 * Lakehouse is licensed under Mulan PSL v2.
 * You can use this software according to the terms and conditions of the Mulan PSL v2.
 * You may obtain a copy of Mulan PSL v2 at:
 *          http://license.coscl.org.cn/MulanPSL2
 * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
 * EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
 * MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
 * See the Mulan PSL v2 for more details.
 */

package com.chinamobile.cmss.lakehouse.service.hive;

import com.chinamobile.cmss.lakehouse.common.exception.BaseException;

import java.sql.Connection;
import java.sql.ResultSet;
import java.sql.Statement;
import java.util.Optional;
import java.util.Set;
import java.util.function.Function;

import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;

@Slf4j
public class HiveSqlExecutor implements AutoCloseable {

    @Autowired
    private Connection connection;

    private HiveConnectionPool hiveConnectionPool;

    private String userName;

    @Autowired
    public HiveSqlExecutor(HiveConnectionPool hiveConnectionPool, String userName) {
        this.hiveConnectionPool = hiveConnectionPool;
        this.userName = userName;
        try {
            connection = hiveConnectionPool.borrowObject(userName);
        } catch (Exception e) {
            throw new RuntimeException(e);
        }
    }

    public void execute(String sql) {
        execute(sql, resultSet -> null);
    }

    public Optional<Set<String>> execute(String sql, Function<ResultSet, Set<String>> extractor) {
        log.info("exec sql: {}", sql);
        try (Statement statement = connection.createStatement()) {
            boolean result = statement.execute(sql);
            if (result && extractor != null) {
                return Optional.ofNullable(extractor.apply(statement.getResultSet()));
            } else {
                return Optional.empty();
            }
        } catch (RuntimeException e) {
            throw e;
        } catch (Exception e) {
            throw new BaseException(e.getMessage());
        }
    }

    @Override
    public void close() throws Exception {
        hiveConnectionPool.returnObject(userName, connection);
    }
}
